# !/usr/bin/env python
# -*-coding:utf-8 -*-

"""
# File       : preprocess.py
# Time       ：2021/12/13 15:07
# Author     ：caoxu
# version    ：python 3.9
# Description：preprocess data
"""

import os
import pandas as pd
import random
import PIL
import logging
from matplotlib.patches import Rectangle
import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings("ignore")
logger = logging.getLogger()
logger.setLevel(logging.INFO)
log_path = './Logs/preprocess.log'
fh = logging.FileHandler(log_path, mode='w')
logger.addHandler(fh)


def create_annots(annots, classes):
    _library_prefix = '/exports/LIBRARY/Library'
    annotations = []

    for i in annots.ImageID.unique():

        full_image_path = annots[annots['ImageID'] == i].iloc[0]['ImagePath']

        # if filename does not exitst, continue
        if not os.path.exists(full_image_path):
            print('ERROR: Image does not exist: {}'.format(full_image_path))
            continue

        labels = []

        # for same image
        for n in range(len(annots[annots['ImageID'] == i])):
            x_min, y_min, x_max, y_max = [i for i in annots[annots['ImageID'] == i].iloc[n]['polygon']]
            width = annots[annots['ImageID'] == i].iloc[n]['width'].astype(int)
            height = annots[annots['ImageID'] == i].iloc[n]['height'].astype(int)

            # check if second x is higher than width or first y is higher than height
            if ((x_max > width) or (y_max > height)):
                print('ERROR: x:{} is higher than widht:{} or y:{} is higher than height:{}'.format(x_max, width,
                                                                                                    y_max, height))
                continue

            labels.append('{},{},{},{},{}'.format(x_min, y_min, x_max, y_max,
                                                  annots[annots['ImageID'] == i].iloc[n]['LabelIndex']))
            # print(_library_prefix+srclocation+'/'+srcname)

        # ommit if no labes for current image
        if len(labels) == 0:
            continue
        # append image with labes to annots
        logger.info('image labels{}'.format(str(full_image_path + ' ' + ' '.join(labels))))
        annotations.append(full_image_path + ' ' + ' '.join(labels))

    return annotations


def save_list_to_file(file_, list_):
    annot_file = open(file_, "w")

    # list in list
    if isinstance(list_[0], list):
        for path, polygon in list_:
            annot_file.write(str(path) + ' ' + str(polygon) + '\n')

    # list
    if isinstance(list_[0], str):
        for full_string in list_:
            annot_file.write(str(full_string) + '\n')

    annot_file.close()


def shuffle_and_split(annots, test=5):
    test = 5  # % for testing
    random.shuffle(annots)
    l = len(annots)
    logger.info('len(annots): '.format(l))
    # print(l)

    training_list = annots[: int(l * ((100 - test) / 100))]
    test_list = annots[int(l * ((100 - test) / 100)): -1]

    logger.info('start save list to train file')
    save_list_to_file('train.txt', training_list)
    logger.info('start save list to test file')
    save_list_to_file('test.txt', test_list)


def draw_rectangels_to_labeld_images(classes, annot_file, save_path, relative_path=False):
    FIG_SIZE = (30, 15)

    annot_file = open(annot_file, "r")
    annotations = annot_file.read().splitlines()

    # every annotation is leafleat
    for l in annotations:

        filename = l.split()[0].split('/')[-1][:-3] + 'png'
        pathfilename = l.split()[0]

        if relative_path:
            save_path = save_path + os.path.dirname(pathfilename)
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        else:
            save_filepath = save_path

        polygons = [int(y) for x in str(l.split('jpg ')[1]).split(',') for y in x.split(' ')]

        im = PIL.Image.open(pathfilename)

        # Display the image
        plt.figure(figsize=FIG_SIZE)
        plt.imshow(im, interpolation='nearest')

        # Get the current reference
        ax = plt.gca()

        # for polygon in [ polygons[i:i+8] for i in range(0,len(polygons),8) ]:
        for pol in [polygons[i:i + 5] for i in range(0, len(polygons), 5)]:
            # pol = [ int(p) for p in pol]
            # print(pol)
            rect = Rectangle((pol[0], pol[1]), int(pol[2] - pol[0]), int(pol[3] - pol[1]), linewidth=3, edgecolor='r',
                             facecolor='none')
            ax.text(pol[0], pol[1], classes[pol[4]], color='red')
            ax.add_patch(rect)

        plt.savefig(save_path + '/' + filename, bbox_inches='tight')
        plt.close('all')
        im.close()


def preprocess_data():
    # y3 working path
    train_dir = 'Y3_openimageV5/'
    code_dir = 'Y3/'
    for p in ['conf', 'images']:
        if not os.path.exists(train_dir+'/'+p):
            os.makedirs(train_dir+'/'+p, exist_ok=True)

    logger.info('start Read csv to dataframe')
    # Read csv to dataframe
    raw_annots = pd.read_csv("Y3_openimageV5/conf/test-annotations-bbox.csv")
    raw_classes = pd.read_csv("Y3_openimageV5/conf/class-descriptions-boxable.csv", names=["name", "label"])

    logger.info('start Create classes series')
    # Create classes series
    classes = raw_classes.label

    logger.info('start Functions for transfering annot data')
    # Functions for transfering annot data
    def get_index_of_class(x):
        return raw_classes[raw_classes.name == str(x)].index[0]

    def get_label_of_class(x):
        return raw_classes[raw_classes.name == str(x)].label.values[0]

    imgs = {}

    def get_width_of_img(img):
        if not img in imgs:
            imgs[img] = PIL.Image.open(img).size
        return imgs[img][0]

    def get_height_of_img(img):
        if not img in imgs:
            imgs[img] = PIL.Image.open(img).size
        return imgs[img][1]

    logger.info('start slice the to lower size for testing')
    # slice the to lower size for testing
    annots = raw_annots[['ImageID', 'XMin', 'YMin', 'XMax', 'YMax', 'LabelName']]  # [:100]

    logger.info('start loc annotations LabelIndex、 label、 Image path')
    annots.loc[:, 'LabelIndex'] = annots.loc[:, 'LabelName'].apply(get_index_of_class)
    annots.loc[:, 'Label'] = annots.loc[:, 'LabelName'].apply(get_label_of_class)
    annots.loc[:, 'ImagePath'] = train_dir + 'images/test/' + annots.loc[:, 'ImageID'] + '.jpg'
    annots.loc[:, 'width'] = annots.loc[:, 'ImagePath'].apply(get_width_of_img)
    annots.loc[:, 'height'] = annots.loc[:, 'ImagePath'].apply(get_height_of_img)

    logger.info('start loc annotations Image XY min max')
    annots.loc[:, 'XMin'] = annots.loc[:, 'XMin'].multiply(annots['width']).astype(int)  # .astype('str')+','
    annots.loc[:, 'XMax'] = annots.loc[:, 'XMax'].multiply(annots['width']).astype(int)  # .astype('str')+','
    annots.loc[:, 'YMin'] = annots.loc[:, 'YMin'].multiply(annots['height']).astype(int)  # .astype('str')+','
    annots.loc[:, 'YMax'] = annots.loc[:, 'YMax'].multiply(annots['height']).astype(int)  # .astype('str')+','
    annots['polygon'] = annots.loc[:, ['XMin', 'YMin', 'XMax', 'YMax']].values.tolist()

    logger.info('start Check if classes are imbalanced')
    # Check if classes are imbalanced
    annots.Label.value_counts().plot.bar(figsize=(35, 20))
    annots.Label.value_counts()

    logger.info('start Create Annotations')
    # Create Annotations
    annotations = create_annots(annots, raw_classes)

    logger.info('start Shuffle and split annotations into train and test parts')
    # Shuffle and split annotations into train and test parts
    shuffle_and_split(annotations, 5)

    #  (optional) Draw labels to images, to check if our labels are sized properly
    # draw_rectangels_to_labeld_images(classes, train_dir + '/conf/test.txt',
    #                                  train_dir + '/images/labeled', relative_path=False)


if __name__ == '__main__':
    preprocess_data()
    print('done')